from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Query
from typing import Optional, List
from datetime import datetime
import json
from core.websocket_manager import ConnectionManager
from schemas.response import ResponseModel
from core.state_manager import state_manager  # 直接导入实例

router = APIRouter()
manager = ConnectionManager()

def json_dumps(obj):
    """JSON编码，支持中文"""
    return json.dumps(obj, ensure_ascii=False)

# 支持路径参数方式
@router.websocket("/ws/{client_id}")
async def websocket_endpoint_path(websocket: WebSocket, client_id: str):
    await handle_websocket_connection(websocket, client_id)

# 支持查询参数方式
@router.websocket("/ws")
async def websocket_endpoint_query(
    websocket: WebSocket,
    client_id: str = Query(..., description="客户端ID")
):
    await handle_websocket_connection(websocket, client_id)

# 抽取公共的处理逻辑
async def handle_websocket_connection(websocket: WebSocket, client_id: str):
    try:
        # 使用共享状态
        state_manager.set(f"ws_client_{client_id}", {
            "connected_at": datetime.now().isoformat(),
            "status": "active"
        })
        
        # 建立连接
        await manager.connect(websocket, client_id)
        
        # 发送欢迎消息，包含消息格式说明
        await websocket.send_text(json_dumps({
            "type": "system",
            "message": "连接成功！消息格式示例：{'type': 'message', 'message': '你的消息'} 或 直接发送文本消息",
            "timestamp": datetime.now().isoformat()
        }))
        
        # 广播新用户加入消息
        await manager.broadcast(
            json_dumps({
                "type": "system",
                "message": f"Client {client_id} joined the chat",
                "timestamp": datetime.now().isoformat()
            })
        )
        
        try:
            while True:
                # 接收消息
                data = await websocket.receive_text()
                print(f"服务端接收到：{data}")
                
                try:
                    # 尝试解析JSON消息
                    try:
                        message_data = json.loads(data)
                        message_type = message_data.get("type", "message")
                        
                        if message_type == "heartbeat":
                            # 处理心跳消息
                            await websocket.send_text(json_dumps({
                                "type": "heartbeat",
                                "timestamp": datetime.now().isoformat()
                            }))
                        else:
                            # 处理普通JSON消息
                            response_message = json_dumps({
                                "type": "message",
                                "client_id": client_id,
                                "message": message_data.get("message", ""),
                                "timestamp": datetime.now().isoformat()
                            })
                            # 先发送给发送者确认
                            await websocket.send_text(json_dumps({
                                "type": "ack",
                                "message": "消息已收到",
                                "original_message": message_data.get("message", ""),
                                "timestamp": datetime.now().isoformat()
                            }))
                            # 然后广播给其他人
                            await manager.broadcast(response_message, exclude=client_id)
                    except json.JSONDecodeError:
                        # 如果不是JSON格式，就当作普通文本消息处理
                        response_message = json_dumps({
                            "type": "message",
                            "client_id": client_id,
                            "message": data,
                            "timestamp": datetime.now().isoformat()
                        })
                        # 先发送给发送者确认
                        await websocket.send_text(json_dumps({
                            "type": "ack",
                            "message": "消息已收到",
                            "original_message": data,
                            "timestamp": datetime.now().isoformat()
                        }))
                        # 然后广播给其他人
                        await manager.broadcast(response_message, exclude=client_id)
                        
                except Exception as e:
                    # 处理其他异常
                    await websocket.send_text(json_dumps({
                        "type": "error",
                        "message": f"消息处理错误: {str(e)}"
                    }))
                    
        except WebSocketDisconnect:
            # 处理连接断开
            await manager.disconnect(client_id)
            await manager.broadcast(
                json_dumps({
                    "type": "system",
                    "message": f"Client {client_id} left the chat",
                    "timestamp": datetime.now().isoformat()
                })
            )
            
    finally:
        # 清理状态
        state_manager.delete(f"ws_client_{client_id}")

# 获取当前连接的客户端列表
@router.get("/ws/clients", response_model=ResponseModel[List[str]])
async def get_clients():
    clients = manager.get_connected_clients()
    return ResponseModel(
        code=200,
        message="获取成功",
        data=clients
    ) 